// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/common/hyperloglog/dense_hll.h>
#include <melon/benchmark.h>
#include <melon/init/init.h>
#include <pollux/common/memory/hash_string_allocator.h>

#define XXH_INLINE_ALL
#include <xxhash.h>

using namespace kumo::pollux;

namespace {

template <typename T>
uint64_t hashOne(T value) {
  return XXH64(&value, sizeof(value), 0);
}

// A benchmark for DenseHll::mergeWith(serialized) API.
//
// Measures the time it takes to merge 2 serialized digests using different
// values for hash bits. Larger values of hash bits corresponds to larger
// digests that are more accurate, but slower to merge. The default number of
// hash bits is 11, while in practice 16 is common.
class DenseHllBenchmark {
 public:
  explicit DenseHllBenchmark(memory::MemoryPool* pool) : pool_(pool) {
    for (auto hashBits : {11, 12, 16}) {
      serializedHlls_[hashBits].push_back(makeSerializedHll(hashBits, 1));
      serializedHlls_[hashBits].push_back(makeSerializedHll(hashBits, 2));
    }
  }

  void run(int hashBits) {
    melon::BenchmarkSuspender suspender;

    HashStringAllocator allocator(pool_);
    common::hll::DenseHll hll(hashBits, &allocator);

    suspender.dismiss();

    for (const auto& serialized : serializedHlls_.at(hashBits)) {
      hll.mergeWith(serialized.data());
    }
  }

 private:
  std::string makeSerializedHll(int hashBits, int32_t step) {
    HashStringAllocator allocator(pool_);
    common::hll::DenseHll hll(hashBits, &allocator);
    for (int32_t i = 0; i < 1'000'000; ++i) {
      auto hash = hashOne(i * step);
      hll.insertHash(hash);
    }
    return serialize(hll);
  }

  static std::string serialize(common::hll::DenseHll& denseHll) {
    auto size = denseHll.serializedSize();
    std::string serialized;
    serialized.resize(size);
    denseHll.serialize(serialized.data());
    return serialized;
  }

  memory::MemoryPool* pool_;

  // List of serialized HLLs to use for merging, keyed by the number of hash
  // bits.
  melon::F14FastMap<int, std::vector<std::string>> serializedHlls_;
};

} // namespace

std::unique_ptr<DenseHllBenchmark> benchmark;

BENCHMARK(mergeSerialized11) {
  benchmark->run(11);
}

BENCHMARK(mergeSerialized12) {
  benchmark->run(12);
}

BENCHMARK(mergeSerialized16) {
  benchmark->run(16);
}

int main(int argc, char** argv) {
  melon::Init init(&argc, &argv);

  memory::MemoryManager::initialize({});
  auto rootPool = memory::memoryManager()->addRootPool();
  auto pool = rootPool->addLeafChild("bm");
  benchmark = std::make_unique<DenseHllBenchmark>(pool.get());

  melon::runBenchmarks();
  return 0;
}
